'''
Created on Jul 18, 2009

@author: mkiyer
'''

import re
import logging

def read_bedgraph(fhd):    
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline.startswith('#'):
            continue
        if thisline.startswith('track'):            
            if re.search(r'type=bedGraph', thisline) is None:
                raise ValueError('file type != bedGraph')
            continue
        if thisline.startswith('browser'):
            # TODO: parse browser line
            continue
        thisfields = thisline.split()
        assert len(thisfields) == 4
        chrom = thisfields[0]
        start = int(thisfields[1])
        end = int(thisfields[2])
        value = float(thisfields[3])
        yield chrom, start, end, value
    
def write_bedgraph(fhd, name, chrom, cov_array):
    write_bedgraph_track(fhd, name)
    write_bedgraph_data(fhd, chrom, cov_array)

def write_bedgraph_track(fhd, name):
    fhd.write("track type=bedGraph name=\"%s\"\n" % (name))    

def write_bedgraph_data(fhd, chrom, cov_array, offset=0):
    if not cov_array:
        return
    # reset state
    cur_v = None
    for i,v in enumerate(cov_array):
        if not cur_v:
            # reset state
            start, cur_v = i, v
            continue
        if v != cur_v:
            # if the interval value is zero, no need to write to the file
            # because the default value is zero
            if cur_v != 0:
                fhd.write("%s\t%d\t%d\t%d\n" % (chrom, start + offset, 
                                                i + offset, cur_v))
            # reset state
            start, cur_v = i, v
    # cleanup and write final interval if necessary
    if cur_v != 0:
        fhd.write("%s\t%d\t%d\t%d\n" % (chrom, start + offset, 
                                        i + 1 + offset, cur_v))

class BEDGraphFileStreamWriter(object):
    
    def __init__(self, fhd, chrom, chrom_start=0):
        self.fhd = fhd
        self.chrom = chrom
        self.chrom_start = chrom_start
        self.chrom_end = self.chrom_start
        self.v = 0
    
    def seek_forward(self, distance):
        self.flush()
        self.chrom_start += distance
        self.chrom_end = self.chrom_start
    
    def write_track_name(self, name):
        self.fhd.write("track type=bedGraph name=\"%s\"\n" % (name))    
    
    def write_stream(self, data):
        if data == None:
            return
        for nextv in data:
            if self.v != nextv:
                if self.v != 0:
                    self.fhd.write("%s\t%d\t%d\t%d\n" %
                                   (self.chrom, self.chrom_start,
                                    self.chrom_end, self.v))
                self.chrom_start = self.chrom_end
                self.v = nextv
            self.chrom_end += 1
    
    def flush(self):
        if self.chrom_start != self.chrom_end:
            # cleanup and write final interval if necessary
            if self.v != 0:
                self.fhd.write("%s\t%d\t%d\t%d\n" % 
                               (self.chrom, self.chrom_start,
                                self.chrom_end, self.v))
            self.chrom_start = self.chrom_end
            self.v = 0